const crypto = require('crypto');

/**
 * 生成AES所需的密钥
 * @returns {Buffer} 返回一个16字节的随机数据作为AES密钥
 */
function generateAesKey() {
  // 生成16字节的随机数据作为AES-128密钥
  return crypto.randomBytes(16);
}

/**
 * 生成随机IV（初始化向量）
 * @param {number} size - IV大小，默认为16字节
 * @returns {Buffer} 返回随机IV
 */
function generateIV(size = 16) {
  return crypto.randomBytes(size);
}

/**
 * AES 加密
 * @param {Buffer} data - 要加密的数据（字节数组）
 * @param {Buffer} key - 密钥（16字节，128位）
 * @param {string} algorithm - 加密算法，默认为 'aes-128-ecb'
 * @returns {Buffer} 加密后的数据（字节数组）
 */
function aesEncrypt(data, key, algorithm = 'aes-128-ecb') {
  // 验证密钥长度
  if (key.length !== 16) {
    throw new Error('密钥长度必须为16字节（128位）');
  }

  let cipher;
  let iv;

  if (algorithm.includes('ecb')) {
    // ECB模式不需要IV
    cipher = crypto.createCipheriv(algorithm, key, null);
  } else if (algorithm.includes('ctr')) {
    // CTR模式需要随机IV（16字节）
    iv = generateIV(16);
    cipher = crypto.createCipheriv(algorithm, key, iv);
  } else {
    // 其他模式需要IV（使用全零IV，实际应用中建议使用随机IV）
    iv = Buffer.alloc(16, 0);
    cipher = crypto.createCipheriv(algorithm, key, iv);
  }

  // 执行加密
  const encrypted = Buffer.concat([cipher.update(data), cipher.final()]);

  // 对于CTR模式，返回IV + 密文的组合
  if (algorithm.includes('ctr') && iv) {
    return Buffer.concat([iv, encrypted]);
  }

  return encrypted;
}

/**
 * AES 解密
 * @param {Buffer} encryptedData - 要解密的数据（字节数组）
 * @param {Buffer} key - 密钥（16字节，128位）
 * @param {string} algorithm - 加密算法，默认为 'aes-128-ecb'
 * @returns {Buffer} 解密后的数据（字节数组）
 */
function aesDecrypt(encryptedData, key, algorithm = 'aes-128-ecb') {
  // 验证密钥长度
  if (key.length !== 16) {
    throw new Error('密钥长度必须为16字节（128位）');
  }

  let decipher;
  let iv;
  let actualCiphertext;

  if (algorithm.includes('ecb')) {
    // ECB模式不需要IV
    decipher = crypto.createDecipheriv(algorithm, key, null);
    actualCiphertext = encryptedData;
  } else if (algorithm.includes('ctr')) {
    // CTR模式：前16字节是IV，后面是密文
    if (encryptedData.length < 16) {
      throw new Error('CTR模式加密数据格式错误：缺少IV');
    }
    iv = encryptedData.slice(0, 16);
    actualCiphertext = encryptedData.slice(16);
    decipher = crypto.createDecipheriv(algorithm, key, iv);
  } else {
    // 其他模式需要IV
    iv = Buffer.alloc(16, 0);
    decipher = crypto.createDecipheriv(algorithm, key, iv);
    actualCiphertext = encryptedData;
  }

  // 执行解密
  const decrypted = Buffer.concat([decipher.update(actualCiphertext), decipher.final()]);

  return decrypted;
}

/**
 * 专门用于CTR模式的加密函数
 * @param {Buffer} data - 要加密的数据
 * @param {Buffer} key - 密钥（16字节）
 * @returns {Buffer} 返回 IV + 密文 的组合
 */
function aesEncryptCTR(data, key) {
  return aesEncrypt(data, key, 'aes-128-ctr');
}

/**
 * 专门用于CTR模式的解密函数
 * @param {Buffer} encryptedData - 要解密的数据（IV + 密文）
 * @param {Buffer} key - 密钥（16字节）
 * @returns {Buffer} 解密后的原始数据
 */
function aesDecryptCTR(encryptedData, key) {
  return aesDecrypt(encryptedData, key, 'aes-128-ctr');
}

// 使用示例
function example() {
  try {
    console.log('=== AES加密解密示例 ===\n');

    // 准备测试数据
    const originalData = Buffer.from('Hello, World! This is a test message for AES-128-CTR encryption.', 'utf8');
    const key = Buffer.from('W1kBIaA0lsQ8KxNx', 'utf8');

    console.log('原始数据:', originalData.toString('utf8'));
    console.log('数据长度:', originalData.length, '字节');
    console.log('密钥:', key.toString('hex'));
    console.log('密钥长度:', key.length, '字节\n');

    // 测试 ECB 模式
    console.log('1. ECB模式测试:');
    const encryptedECB = aesEncrypt(originalData, key, 'aes-128-ecb');
    const decryptedECB = aesDecrypt(encryptedECB, key, 'aes-128-ecb');
    console.log('ECB加密后长度:', encryptedECB.length, '字节');
    console.log('ECB解密验证:', originalData.equals(decryptedECB) ? '成功' : '失败');
    console.log('ECB解密文本:', decryptedECB.toString('utf8'));

    // 测试 CTR 模式
    console.log('2. CTR模式测试:');
    const encryptedCTR = aesEncrypt(originalData, key, 'aes-128-ctr');
    const decryptedCTR = aesDecrypt(encryptedCTR, key, 'aes-128-ctr');
    console.log('CTR加密后长度:', encryptedCTR.length, '字节（包含16字节IV）');
    console.log('实际密文长度:', encryptedCTR.length - 16, '字节');
    console.log('CTR解密验证:', originalData.equals(decryptedCTR) ? '成功' : '失败');
    console.log('CTR解密文本:', decryptedCTR.toString('utf8'));

    // 测试专门的CTR函数
    console.log('3. 专用CTR函数测试:');
    const encryptedCTR2 = aesEncryptCTR(originalData, key);
    const ctrDecryptHex = encryptedCTR2.toString("hex");
    const ctrDecryptBase64 = encryptedCTR2.toString("base64");
    console.log('专用CTR加密结果Hex:', ctrDecryptHex);
    console.log('专用CTR加密结果Base64:', ctrDecryptBase64);
    const decryptedCTR2 = aesDecryptCTR(encryptedCTR2, key);
    console.log('专用CTR解密验证:', originalData.equals(decryptedCTR2) ? '成功' : '失败\n');

    // 测试不同长度的数据
    console.log('4. 不同长度数据测试:');
    const shortData = Buffer.from('Hi', 'utf8');
    const longData = Buffer.from('A'.repeat(100), 'utf8');

    const shortEncrypted = aesEncryptCTR(shortData, key);
    const shortDecrypted = aesDecryptCTR(shortEncrypted, key);
    console.log('短数据(2字节)加解密:', shortData.equals(shortDecrypted) ? '成功' : '失败');

    const longEncrypted = aesEncryptCTR(longData, key);
    const longDecrypted = aesDecryptCTR(longEncrypted, key);
    console.log('长数据(100字节)加解密:', longData.equals(longDecrypted) ? '成功' : '失败');

  } catch (error) {
    console.error('错误:', error.message);
    console.error(error.stack);
  }
}

// 导出函数供其他模块使用
module.exports = {
  generateAesKey,
  generateIV,
  aesEncrypt,
  aesDecrypt,
  aesEncryptCTR,
  aesDecryptCTR
};

// 如果直接运行此文件，执行示例
if (require.main === module) {
  example();
}
